#pragma once

#include <iostream>
#include <vector>

using namespace std;


class UnionFindSet
{
private:
	vector<int> _union_set;
public:
	UnionFindSet(size_t n) {
		_union_set.resize(n, -1);
	}

	UnionFindSet(vector<int>& it) 
		: _union_set(it)
	{
	}

	// 将两个下标集合
	void _union(size_t a, size_t b)
	{
        a = GetRoot(a);
        b = GetRoot(b);
        if (a == b) return; // 相同树，不合并

		// 让小树合并到大树
		if (abs(_union_set[a]) > abs(_union_set[b]))
			swap(a, b);

		_union_set[a] += _union_set[b];
		_union_set[b] = a;
        //可以添加优化树结构操作，缩小树高度
	}

	size_t GetRoot(size_t order)
	{
		int root = order;
		while (_union_set[root] >= 0)
		{
			root = _union_set[root];
		}
		
		// 开始向上路径压缩
		while (_union_set[order] >= 0)
		{
			int parent = _union_set[order];
			_union_set[order] = root;
			order = parent;
		}
		return root;
	}

	size_t size()
	{
		size_t num = 0;
		for (auto& it : _union_set)
		{
			if (it < 0) num++;
		}

        for (auto i : _union_set)
        {
            std::cout << i  << " ";
        }
        cout << endl;
		return num;
	}
};